"""
Copyright 2020 The OneFlow Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import oneflow
from oneflow.framework.docstr.utils import add_docstr

add_docstr(
    oneflow.read_onerec,
    r"""
    read_onerec(files:List[str], batch_size:int, random_shuffle:bool, shuffle_mode:str, shuffle_buffer_size=1024, shuffle_after_epoch=False, verify_example=True, placement=None, sbp=None) -> Tensor 
    
    Read OneRec format dataset into a Tensor which then can be decode by decode_onerec API.

    Args:      
        files: (List[str]): The file list to be read from filesystem       
        batch_size(int): batch size
        random_shuffle(bool): shuffle or not
        shuffle_mode(str): can be "batch" or "instance"
        shuffle_buffer_size(int): shuffle buffer size, default to 1024
        shuffle_after_epoch(bool): if shuffle after each epoch
        verify_example(bool): if verify example, defaults to True
        placement(Optional[oneflow._oneflow_internal.placement]): The placement attribute allows you to specify which physical device the tensor is stored on.
        sbp(Optional[Union[oneflow._oneflow_internal.sbp.sbp, List[oneflow._oneflow_internal.sbp.sbp]]]): When creating a consistent tensor, specify the SBP of the tensor.
    
        
    For example:

    .. code-block:: python

        import oneflow as flow
        files = ['file01.onerec', 'file02.onerec']
        readdata_1 = flow.read_onerec(files, 10, True, "batch")
      
        # decode readdata_1 ... 

    .. code-block:: python

        import oneflow as flow
        files = ['file01.onerec', 'file02.onerec']
        readdata_2 = flow.read_onerec(
            files,
            batch_size=10,
            random_shuffle=True,
            shuffle_mode="batch",
            placement=flow.placement("cpu", {0: [0]}),
            sbp=[flow.sbp.split(0)],
        )    
      
        # decode readdata_2 ... 

    """,
)

add_docstr(
    oneflow.decode_onerec,
    r"""
    (Tensor input, String key, DataType dtype, Shape shape, Bool is_dynamic=False, Shape reshape=None, Shape batch_padding=None) 
    decode_onerec(input:Tensor, key:str, dtype, shape, is_dynamic=False, reshape=None, batch_padding=None) -> Tensor 
    
    Decode a tensor from input which should be generated before by oneflow.read_onerec.

    Args:      
        input: (Tensor): The tensor generated by oneflow.read_onerec before.
        key(str): The field name of the tensor to be decode
        shape(bool):  The shape of the tensor to be decode
        is_dynamic(bool): The tensor shape is dynamic or not
        reshape(tuple): Set it if you want to reshape the tensor
        batch_padding(tuple): Set it if batch padding is needed
    
        
    For example:

    .. code-block:: python

        import oneflow as flow
        files = ['file01.onerec', 'file02.onerec']
        # read onerec dataset form files
        readdata = flow.read_onerec(files, 10, True, "batch")
      
        # decode
        labels = flow.decode_onerec(readdata, key="labels", dtype=flow.int32, shape=(1,))
        dense_fields = flow.decode_onerec(readdata, key="dense_fields", dtype=flow.float, shape=(13,))

    """,
)
